package kvraft

import (
	"bytes"
	"go-kvraft/labgob"
	"go-kvraft/raft"
	"go-kvraft/rpcutil"
	"log"
	"net/rpc"
	"sync"
	"sync/atomic"
	"time"
)

const Debug = true

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug {
		log.Printf(format, a...)
	}
	return
}

type KVServer struct {
	mu           sync.RWMutex       // 读写锁，用来保护共享资源的读写操作
	me           int                // 该 KVServer 实例的编号
	rf           *raft.Raft         // 该 KVServer 实例所对应的 Raft 实例
	applyCh      chan raft.ApplyMsg // 用来接收 Raft 模块提交的日志记录
	dead         int32              // set by Kill() 用来表示该 KVServer 实例是否被 Kill
	maxraftstate int                //  用来设置当 Raft 日志的大小达到该值时，需要进行持久化操作
	// Your definitions here.
	persister   *raft.Persister        // 持久化模块，用来将 Raft 模块的状态持久化到磁盘中
	state       map[string]string      // 用来存储 key-value 数据
	client      map[int64]int          // 用来存储每个 client 的最近一次请求的序列号
	broadcastCh map[int]chan Operation // 用来存储每个 client 对应的 channel，用于通知 client 的请求结果
}

// Operation 操作的相关信息
type Operation struct {
	ClientId  int64 // 客户端的唯一标识符，用于区分不同的客户端
	MessageId int   // 消息的唯一标识符，用于防止重复请求
	Key       string
	Value     string
	Method    string // 操作的类型，是 Put 还是 Append
}

// Snapshot
// 用来保存当前 key-value 存储系统的状态
// 以及 client 端发送的最新请求的信息，以便在需要时进行快照和恢复
type Snapshot struct {
	State  map[string]string // 用于存储当前 key-value 存储系统中所有的 key-value 对
	Client map[int64]int     // 用于记录每个 client 最近一次请求的信息，包括 client id 和 message id
}

// KVServer 中用于做快照的函数，
// 其输入参数 index 表示当前需要做快照的索引值
func (kv *KVServer) snapshot(index int) {
	// 造一个 Snapshot 结构体，将 KVServer 中的状态 state 和 client 存入其中
	snapshot := Snapshot{
		State:  kv.state,
		Client: kv.client,
	}
	w := new(bytes.Buffer)
	// 将 snapshot 对象编码后，将其写入缓存中
	e := labgob.NewEncoder(w)
	e.Encode(snapshot)
	// 将快照信息和索引值传递给底层的 raft 模块进行存储
	kv.rf.Snapshot(index, w.Bytes())
}

// 这个函数是用来读取kv服务器状态的快照的
func (kv *KVServer) readPersist(snapshot []byte) {
	if snapshot == nil || len(snapshot) < 1 {
		return
	}

	r := bytes.NewBuffer(snapshot)
	d := labgob.NewDecoder(r)

	var decodeSnapshot Snapshot
	if err := d.Decode(&decodeSnapshot); err == nil {
		kv.state = decodeSnapshot.State
		kv.client = decodeSnapshot.Client
	} else {
		panic(err)
	}
}

// Kill 函数用于关闭 KVServer 实例
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1) // 将 dead 标志位设置为 1，表示该 KVServer 实例已经关闭
	kv.rf.Kill()                   // 关闭对应的 Raft 实例
	// Your code here, if desired.
}

// killed函数用于判断当前KVServer实例是否已经关闭
func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead) // 读取dead标志位的值
	return z == 1                   // 如果dead标志位的值为1，则返回true，否则返回false
}

// applyRoutine函数是一个协程，用于处理applyCh中传入的applyMsg消息
func (kv *KVServer) applyRoutine() {
	for !kv.killed() {
		applyMsg := <-kv.applyCh              // 从applyCh中获取applyMsg消息
		kv.mu.Lock()                          // 加锁，以确保状态和客户端映射的同步更新
		commandValid := applyMsg.CommandValid // 检查该applyMsg消息是否包含Command
		if commandValid {                     // 检查该applyMsg消息是否包含Command
			command := applyMsg.Command.(Operation) // 将Command转换为Operation类型
			commandIndex := applyMsg.CommandIndex
			messageId := command.MessageId
			clientId := command.ClientId
			method := command.Method
			key := command.Key
			value := command.Value
			// 如果该消息的MessageId大于该客户端的最新MessageId，则更新状态
			if kv.client[clientId] < messageId {
				switch method {
				case "Put":
					kv.state[key] = value
				case "Append":
					kv.state[key] += value
				}
				kv.client[clientId] = messageId
			}

			// // 如果有监听该Command Index的Channel，则向该Channel发送Command消息
			if broadcastCh, ok := kv.broadcastCh[commandIndex]; ok {
				broadcastCh <- command
			}

			// 如果达到了快照条件（kv.maxraftstate不等于-1且当前raft状态大小大于kv.maxraftstate），则进行快照
			if kv.maxraftstate != -1 && kv.persister.RaftStateSize() > kv.maxraftstate {
				kv.snapshot(commandIndex)
			}
		}
		snapshotValid := applyMsg.SnapshotValid
		// 检查该applyMsg消息是否包含Snapshot
		if snapshotValid {
			// 将Snapshot恢复为decodeSnapshot
			snapshot := applyMsg.Snapshot
			snapshotTerm := applyMsg.SnapshotTerm
			snapshotIndex := applyMsg.SnapshotIndex
			// 如果当前Term大于等于Snapshot的Term，并且当前Index大于等于Snapshot的Index，则进行快照安装
			if kv.rf.CondInstallSnapshot(snapshotTerm, snapshotIndex, snapshot) {
				kv.readPersist(snapshot)
			}
		}
		kv.mu.Unlock()
	}
}

// Operation 函数是用于处理客户端请求的函数
func (kv *KVServer) Operation(request *OperationRequest, response *OperationResponse) error {
	clientId := request.ClientId
	messageId := request.MessageId
	key := request.Key
	value := request.Value
	method := request.Method
	command := Operation{
		ClientId:  clientId,
		MessageId: messageId,
		Key:       key,
		Value:     value,
		Method:    method,
	}
	index, _, isLeader := kv.rf.Start(command) // 将command发送给Raft模块
	if !isLeader {                             // 如果当前KVServer实例不是Leader，则返回错误信息
		response.Err = ErrWrongLeader
		return nil
	}
	kv.mu.Lock()
	kv.broadcastCh[index] = make(chan Operation, 1) // 创建一个channel，用于接收该command的结果
	broadcastCh := kv.broadcastCh[index]
	kv.mu.Unlock()

	select {
	case appliedCommand := <-broadcastCh: // 等待Raft模块将该command应用到状态机上
		kv.mu.RLock()
		if method == "Get" {
			response.Value = kv.state[key] // 如果该command是Get操作，则将对应的value返回给客户端
		}
		kv.mu.RUnlock()

		appliedClientId := appliedCommand.ClientId
		appliedMessageId := appliedCommand.MessageId
		if clientId != appliedClientId || messageId != appliedMessageId { // 检查返回的command是否与发送的command一致
			response.Err = ErrWrongLeader
			return nil
		}
		response.Err = OK // 如果一致，则返回OK

	case <-time.After(500 * time.Millisecond): // 如果等待超时，则返回错误信息
		response.Err = ErrTimeout
		return nil
	}

	go func() { // 关闭channel并删除对应的映射
		kv.mu.Lock()
		close(broadcastCh)
		delete(kv.broadcastCh, index)
		kv.mu.Unlock()
	}()
	return nil
}

// StartKVServer 函数用于启动一个KVServer实例
// servers[]是一个包含所有KVServer实例的地址的数组
// me是当前KVServer实例在servers[]中的索引
// persister是一个持久化模块，用于将Raft模块的状态持久化到磁盘中
// maxraftstate用于设置当Raft日志的大小达到该值时，需要进行持久化操作
func StartKVServer(servers []*rpcutil.ClientEnd, me int, persister *raft.Persister, maxraftstate int) *KVServer {
	// 在Go的RPC库中注册需要进行编码和解码的结构体
	labgob.Register(Operation{})
	labgob.Register(Snapshot{})

	// 创建一个KVServer实例
	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate

	// 初始化KVServer实例中的状态
	kv.state = make(map[string]string)
	kv.client = make(map[int64]int)
	kv.broadcastCh = make(map[int]chan Operation)

	// 将KVServer实例中的状态从快照中恢复
	kv.persister = persister
	kv.readPersist(kv.persister.ReadSnapshot())

	// 创建一个用于接收Raft模块提交的日志记录的channel
	kv.applyCh = make(chan raft.ApplyMsg)
	// 创建一个Raft实例
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	if err := rpc.Register(kv.rf); err != nil {
		panic(err)
	}

	// 启动一个协程，用于处理applyCh中传入的applyMsg消息
	go kv.applyRoutine()
	return kv

}
